package com.toymate.ai.llm.model.dify;

import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.toymate.ai.llm.AbstractLlmService;
import com.toymate.ai.llm.LlmConfig;
import com.toymate.modules.iot.entity.IotRole;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * DifyLlmService
 *
 * @author weirdor (179946422@qq.com)
 * @version v1.0.0
 */
@Slf4j
public class DifyLlmService extends AbstractLlmService {

    private OkHttpClient okHttpClient;

    public DifyLlmService(LlmConfig config, IotRole role) {
        this.role = role;
        this.config = config;
        this.okHttpClient = new OkHttpClient.Builder()
                .connectTimeout(1, TimeUnit.DAYS)
                .readTimeout(1, TimeUnit.DAYS)
                .build();
    }

    @Override
    public void chat(String query, String productId, String deviceNum) {
        chat(query, RandomUtil.randomString(10), IdUtil.getSnowflakeNextIdStr(), productId, deviceNum);
    }

    @Override
    public void chat(String query, String userId, String productId, String deviceNum) {
        chat(query, userId, IdUtil.getSnowflakeNextIdStr(), productId, deviceNum);
    }

    @Override
    public void chat(String query, String userId, String conversationId, String productId, String deviceNum) {
        sendChatRequest(query, userId, conversationId, productId, deviceNum);
    }

    private void sendChatRequest(String query, String userId, String conversationId, String productId, String deviceNum) {
        Map<String, Object> inputs = new HashMap<>();
        log.info("提示词:{}", role.getPrompt());
        inputs.put("prompt", role.getPrompt());
        Map<String, Object> map = MapUtil.newHashMap();
        map.put("inputs", inputs);
        map.put("query", query);
        map.put("response_mode", "streaming");
        if (StrUtil.isEmpty(conversationId)) {
            map.put("conversation_id", "");
        }else {
            map.put("conversation_id", conversationId);
        }
        map.put("user", 1);
        String requestJson = JSONUtil.toJsonStr(map);
        log.info("请求参数:{}", requestJson);
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, requestJson);
        Request request = new Request.Builder()
                .url(config.getDifyConfig().getApiBase() + "/chat-messages")
                .addHeader("Authorization", "Bearer " + config.getDifyConfig().getKey())
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "*/*")
                .addHeader("Connection", "keep-alive")
                .post(body)
                .build();
        request = request.newBuilder().post(body).build();
        EventSource.Factory factory = EventSources.createFactory(okHttpClient);
        EventSourceListener sourceListener = new DifySSEEventSourceListener(productId, deviceNum, query, role);
        //创建事件
        factory.newEventSource(request, sourceListener);
    }
}
